﻿#region LicenseHeader

// Copyright 2012 The Trustees of Indiana University.  All rights reserved.
// 
// The Indiana University licenses this file to you under the Apache License, 
// Version 2.0 (the "License"); you may not use this file except in compliance 
// with the License.  You may obtain a copy of the License at 
// 
// http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software 
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
//  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  
// See the License for the specific language governing permissions and 
// limitations under the License.
// 
// @author Thilina Gunarathne (tgunarat@indiana.edu)

#endregion

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using AzureMRCore.DataTypes;
using AzureMRCore.InputFormat;
using AzureMRCore.MapRed;
using Microsoft.WindowsAzure;
using Microsoft.WindowsAzure.ServiceRuntime;
using Microsoft.WindowsAzure.StorageClient;

namespace AzureMRCore.IterativeDataCache
{
    /*
     * Make this a seperate thread... managing the storage limits of the data cache
     */

    public class DataCache
    {
        private readonly CloudBlobClient _blobClient;
        public ConcurrentDictionary<string, CachedItem> InMemCache;

        private readonly Dictionary<string, string> _diskCache;
        private readonly List<string> _downloadingBlobs;

        public DataCache(CloudStorageAccount storageAccount)
        {
            _blobClient = storageAccount.CreateCloudBlobClient();
            InMemCache = new ConcurrentDictionary<string, CachedItem>();
            _diskCache = new Dictionary<string, string>();
            _downloadingBlobs = new List<string>();
        }

        //up to the formatter to load it to memory or not. implement WordCout interations..
        /* TODO: temporararily adding JobID  restricting caching per job due to table API limitations */

        public IInputFormat<K, V> GetData<K, V>(string blobURI, IInputFormat<K, V> inputFormat, string jobID,
                                                string iteration, string appName,
                                                string workerID, string mapID, bool isBCast = false)
            where K : Key
            where V : IValue
        {
            if (InMemCache.ContainsKey(blobURI))
            {
                //TODO: use tryGet
                var cachedItem = InMemCache[blobURI];
                inputFormat = (IInputFormat<K, V>) (cachedItem.CachedInputFormat.GetCachedCopy());
                cachedItem.MapID = mapID;
                cachedItem.AppName = appName;
                cachedItem.Iteration = Int32.Parse(iteration);
                cachedItem.IsBCastData = isBCast;
                cachedItem.JobID = jobID;
                //Trace.WriteLine("Got from Memory Cache in DataCache.cs.");
            }
            else
            {
                //Stopwatch sw = Stopwatch.StartNew();
                if (inputFormat is CachedShortMatrixFormat)
                {
                    var shortMatrixFormat = (CachedShortMatrixFormat) inputFormat;
                    shortMatrixFormat.Init(blobURI, _blobClient);
                }
                else
                {
                    string localFileName = DownloadFile(blobURI, jobID, workerID);
                    inputFormat.Init(localFileName);
                }
                //sw.Stop();
                //Trace.WriteLine("Download & init time:" + sw.ElapsedMilliseconds + "  " + blobURI);

                if (inputFormat is ICachedInputFormat)
                {
                    lock (InMemCache)
                    {
                        if (!(InMemCache.ContainsKey(blobURI)))
                        {
                            ICachedInputFormat cachedInputFormat = (ICachedInputFormat) inputFormat;
                            CachedItem cachedItem = new CachedItem(cachedInputFormat, blobURI, jobID,
                                                                   Int32.Parse(iteration), appName,
                                                                   cachedInputFormat.GetSize())
                                                        {
                                                            MapID = mapID,
                                                            IsBCastData = isBCast
                                                        };
                            InMemCache.TryAdd(blobURI, cachedItem);
                        }
                    }
                }
            }
            return inputFormat;
        }

        private string DownloadFile(string blobURI, string jobID, string workerID)
        {
            bool isDownloading = false;

            string file;
            if (_diskCache.TryGetValue(blobURI, out file))
            {
                //Trace.WriteLine("Got from Disk Cache directly.");
                return file;
            }

            lock (_downloadingBlobs)
            {
                if (_downloadingBlobs.Contains(blobURI))
                {
                    isDownloading = true;
                }
                else
                {
                    _downloadingBlobs.Add(blobURI);
                }
            }

            if (isDownloading)
            {
                while (_downloadingBlobs.Contains(blobURI))
                {
                    Thread.Sleep(10);
                }
//                Trace.WriteLine("Got from Disk Cache by waiting");
                return _diskCache[blobURI];
            }

            //adding the try catch to make sure the blob gets removed from downloading blobs in any case.
            try
            {
                LocalResource localResource = RoleEnvironment.GetLocalResource("temp");
                string downloadDir = localResource.RootPath + "\\" + workerID + "\\" + Guid.NewGuid();

                file = MapRedUtils.DownloadBLOB(downloadDir, _blobClient, jobID, blobURI);
                _diskCache.Add(blobURI, file);
            }
            catch (Exception e)
            {
                throw e;
            }
            finally
            {
                lock (_downloadingBlobs)
                {
                    _downloadingBlobs.Remove(blobURI);
                }
            }
            return file;
        }

        //cleanup broadcast data for iterations before the given
        public void CleanupBCastData(string jobID, string appName, int iteration)
        {
            IEnumerable<string> oldDataKeys =
                InMemCache.Where(
                    entry =>
                    ((entry.Value.JobID == jobID) & (entry.Value.AppName == appName) &
                     (entry.Value.Iteration < iteration) & (entry.Value.IsBCastData))).Select(entry => entry.Key);
            List<string> oldDataKeysList = oldDataKeys.ToList();
            foreach (string key in oldDataKeysList)
            {
                CachedItem temp;
                InMemCache.TryRemove(key, out temp);
            }
        }
    }
}